package com.liuqi.openai.core.merge;

import com.liuqi.openai.core.chat.ChatCompletionResponse;
import com.liuqi.openai.core.chat.ToolCall;
import com.liuqi.openai.core.completion.CompletionResponse;
import com.liuqi.openai.core.shared.Usage;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 流式响应数据库合并
 *
 * @author liuqi
 * @date 2025/8/6
 **/
public abstract class MergeChunk<Response> {

    // AI 公共的响应数据

    protected String id;
    protected Integer created;
    protected String model;
    protected String object;
    protected Usage usage;

    // 因为 AI 响应 choices 是数组, 虽然大部分情况下的只有一条数据, 但也为了防止在未来 choices 会有多条数据

    /**
     * AI 正文内容
     *
     * key = index, 与 choices 的 index一一对应
     */
    protected Map<Integer, StringBuilder> contentBuff;

    /**
     * AI finish_reason
     *
     * key = index, 与 choices 的 index一一对应
     */
    protected Map<Integer, StringBuilder> finishReasonBuff;

    /**
     * 角色, 一般都是: assistant
     *
     * key = index, 与 choices 的 index一一对应
     */
    protected Map<Integer, StringBuilder> roleBuff;

    /**
     * AI 思考内容
     *
     * key = index, 与 choices 的 index一一对应
     */
    protected Map<Integer, StringBuilder> reasoningContentBuff;

    /**
     * AI 工具调用
     *
     * key = index, 与 choices 的 index一一对应
     */
    protected Map<Integer, List<ToolCall>> toolCallsBuff;

    /**
     * 根据响应数据类型获取解析数据块的对象
     *
     * @param responseClass
     * @return
     */
    public static MergeChunk getMergeChunk(Class<?> responseClass) {
        if (responseClass == ChatCompletionResponse.class) {
            return new MergeChunkChatCompletion();
        }
        if (responseClass == CompletionResponse.class) {
            return new MergeChunkCompletion();
        }
        throw new IllegalArgumentException("未定义此类型合并数据块的工具: responseClass = " + responseClass);
    }

    /**
     * 合并AI流式响应块
     *
     * @param response  AI流式响应块
     */
    public abstract void merge(Response response);

    /**
     * 获取合并AI流式响应块的结果数据
     *
     */
    public abstract Response getResult();

    protected void setCommon(String id,
                             Integer created,
                             String model,
                             String object,
                             Usage usage) {
        if (this.id == null && id != null ) {
            this.id = id;
        }
        if (this.created == null && created != null) {
            this.created = created;
        }
        if (this.model == null && model != null) {
            this.model = model;
        }
        if (this.object == null && object != null) {
            this.object = object;
        }
        if (usage != null) {
            this.usage = usage;
        }
    }

    /**
     * 添加AI响应内容
     *
     * @param index
     * @param content
     */
    protected void appendContent(int index, String content) {
        if (contentBuff == null) {
            contentBuff = new HashMap<>(1);
        }
        append(index, content, contentBuff);
    }

    /**
     * 添加AI思考内容
     *
     * @param index
     * @param reasoningContent
     */
    protected void appendReasoningContent(int index, String reasoningContent) {
        if (reasoningContentBuff == null) {
            reasoningContentBuff = new HashMap<>(1);
        }
        append(index, reasoningContent, reasoningContentBuff);
    }

    /**
     * 添加角色
     *
     * @param index
     * @param role
     */
    protected void appendRole(int index, String role) {
        if (roleBuff == null) {
            roleBuff = new HashMap<>(1);
        }
        append(index, role, roleBuff);
    }

    /**
     * 添加工具调用
     *
     * @param index
     * @param toolCall
     */
    protected void appendToolCalls(int index, ToolCall toolCall) {
        if (toolCall == null) {
            return;
        }
        if (toolCallsBuff == null) {
            toolCallsBuff = new HashMap<>(1);
        }
        List<ToolCall> toolCalls = toolCallsBuff.get(index);
        if (toolCalls == null) {
            toolCallsBuff.put(index, toolCalls = new ArrayList<>(6));
        }
        toolCalls.add(toolCall);
    }

    /**
     * 设置 finish_reason
     *
     * @param index
     * @param finishReason
     */
    protected void setFinishReason(int index, String finishReason) {
        if (finishReasonBuff == null) {
            finishReasonBuff = new HashMap<>(1);
        }
        append(index, finishReason, finishReasonBuff);
    }

    /**
     * 获取思考内容
     *
     * @param index 索引
     */
    protected String getReasoningContent(int index) {
        StringBuilder reasoning;
        if (reasoningContentBuff == null || (reasoning = reasoningContentBuff.get(index)) == null) {
            return null;
        }
        return reasoning.toString();
    }

    /**
     * 获取AI响应内容
     *
     * @param index 索引
     */
    protected String getContent(int index) {
        StringBuilder content;
        if (contentBuff == null || (content = contentBuff.get(index)) == null) {
            return null;
        }
        return content.toString();
    }

    /**
     * 获取 finish_reason
     *
     * @param index 索引
     */
    protected String getFinishReason(int index) {
        StringBuilder finishReason;
        if (finishReasonBuff == null || (finishReason = finishReasonBuff.get(index)) == null) {
            return null;
        }
        return finishReason.toString();
    }

    /**
     * 获取角色
     *
     * @param index 索引
     */
    protected String getRole(int index) {
        StringBuilder role;
        if (roleBuff == null || (role = roleBuff.get(index)) == null) {
            return null;
        }
        return role.toString();
    }

    private void append(int index, String str, Map<Integer, StringBuilder> mapBuf) {
        if (str == null) return;
        StringBuilder buf = mapBuf.get(index);
        if (buf == null) {
            mapBuf.put(index, buf = new StringBuilder());
        }
        buf.append(str);
    }

}
